#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 'utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

def writeIpv4ForwardRule(p4info_helper, sw, dst_ip, dst_mac, port):
    """Install IPv4 forwarding rule."""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip, 32)
        },
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac,
            "port": port
        }
    )
    sw.WriteTableEntry(table_entry)
    print("Installed IPv4 forward rule for %s -> %s" % (dst_ip, dst_mac))

def writeAclRule(p4info_helper, sw, dst_ip, udp_port):
    """Install ACL rule to drop packets."""
    match_fields = {}

    if dst_ip:
        match_fields["hdr.ipv4.dstAddr"] = dst_ip
    if udp_port is not None:
        match_fields["hdr.udp.dstPort"] = udp_port

    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl",
        match_fields=match_fields,
        action_name="MyIngress.drop"
    )
    sw.WriteTableEntry(table_entry)
    print("Installed ACL rule: IP=%s, UDP_port=%s" % (dst_ip, udp_port))

def main(p4info_file_path, bmv2_file_path):
    # Load P4 program info
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Connect to switch s1 via P4Runtime gRPC
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')

        # Establish controller as master
        s1.MasterArbitrationUpdate()
        # Push P4 program to the switch
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                      bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 program using SetForwardingPipelineConfig")

        # Add IPv4 forwarding rules
        writeIpv4ForwardRule(p4info_helper, s1, "10.0.1.1", "00:00:00:00:01:01", 1)
        writeIpv4ForwardRule(p4info_helper, s1, "10.0.1.2", "00:00:00:00:01:02", 2)
        writeIpv4ForwardRule(p4info_helper, s1, "10.0.1.3", "00:00:00:00:01:03", 3)
        writeIpv4ForwardRule(p4info_helper, s1, "10.0.1.4", "00:00:00:00:01:04", 4)

        # Install ACL rules as specified in requirements
        # Rule 1: Drop all UDP packets with dstPort=80
        writeAclRule(p4info_helper, s1, None, 80)

        # Rule 2: Drop all packets with dstIP=10.0.1.4
        writeAclRule(p4info_helper, s1, "10.0.1.4", None)

        print("Successfully installed all rules!")

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action='store', default='./build/acl.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action='store', default='./build/acl.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        sys.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        sys.exit(1)

    main(args.p4info, args.bmv2_json)
